from ultralytics import YOLO
import cv2
import os
import argparse
import numpy as np
from datetime import datetime

def parse_args():
    parser = argparse.ArgumentParser(description='Test YOLOv8 Fruit Recognition Model')
    parser.add_argument('--model', type=str, default='models/best.pt', help='Model path')
    parser.add_argument('--source', type=str, default='data/test/images', help='Test images directory')
    parser.add_argument('--conf', type=float, default=0.25, help='Confidence threshold')
    parser.add_argument('--device', type=str, default='', help='Device to run on')
    parser.add_argument('--output', type=str, default='validation_results.txt', help='Output file for validation results')
    return parser.parse_args()

def main():
    args = parse_args()
    
    # Load the model
    model = YOLO(args.model)
    
    # Run inference on test images
    results = model(args.source, conf=args.conf, device=args.device)
    
    # Create output directory if it doesn't exist
    output_dir = 'results'
    os.makedirs(output_dir, exist_ok=True)
    
    # Prepare results file
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    results_file = os.path.join(output_dir, args.output)
    
    # Statistics variables
    total_images = len(results)
    correct_detections = 0
    confidence_sum = 0
    class_stats = {}
    
    # Start writing to the file
    with open(results_file, 'w', encoding='utf-8') as f:
        f.write(f"===== YOLOv8 Fruit Recognition Model Validation Results =====\n")
        f.write(f"Date and Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"Model: {args.model}\n")
        f.write(f"Confidence Threshold: {args.conf}\n")
        f.write(f"Test Images: {args.source}\n")
        f.write(f"Total Images: {total_images}\n\n")
        f.write("===== Detailed Detection Results =====\n\n")
        
        # Process results
        for i, result in enumerate(results):
            # Get image filename from the path in the result object
            img_path = result.path
            img_filename = os.path.basename(img_path)
            
            # Try to determine true class from filename
            # Assuming filenames are in format like "apple_123.jpg" or "banana_r_456.jpg"
            true_class = None
            for fruit in ['apple', 'banana', 'grape', 'kiwi', 'lemon', 'mango', 'orange', 'pineapple', 'strawberry', 'watermelon']:
                if fruit in img_filename.lower():
                    true_class = fruit
                    break
            
            # Save the result image with detections
            result_img = result.plot()
            cv2.imwrite(os.path.join(output_dir, f'result_{i}.jpg'), result_img)
            
            # Write detection information to file
            f.write(f"Image {i}: {img_filename}\n")
            if true_class:
                f.write(f"  True Class: {true_class}\n")
            
            # Initialize detected flag and highest confidence for this image
            detected_class = None
            highest_conf = 0
            
            boxes = result.boxes
            for box in boxes:
                cls_id = int(box.cls.item())
                conf = box.conf.item()
                cls_name = model.names[cls_id]
                f.write(f"  Detected {cls_name} with confidence {conf:.2f}\n")
                
                # Keep track of the class with highest confidence
                if conf > highest_conf:
                    highest_conf = conf
                    detected_class = cls_name
            
            # Update statistics if we have true class information
            if true_class and detected_class:
                # Add to class statistics
                if true_class not in class_stats:
                    class_stats[true_class] = {'correct': 0, 'total': 0, 'confidence_sum': 0}
                
                class_stats[true_class]['total'] += 1
                
                # Check if detection was correct
                if true_class == detected_class:
                    correct_detections += 1
                    class_stats[true_class]['correct'] += 1
                    class_stats[true_class]['confidence_sum'] += highest_conf
                
                # Add assessment of correctness to results
                f.write(f"  Assessment: ")
                if true_class == detected_class:
                    f.write(f"CORRECT [正确]\n")
                else:
                    f.write(f"INCORRECT [错误] - True: {true_class}, Detected: {detected_class}\n")
                
                confidence_sum += highest_conf
            
            f.write("\n")
        
        # Write summary statistics
        f.write("===== Summary Statistics =====\n\n")
        
        if total_images > 0 and correct_detections > 0:
            accuracy = correct_detections / total_images * 100
            avg_confidence = confidence_sum / total_images
            
            f.write(f"Overall Accuracy: {accuracy:.2f}% ({correct_detections}/{total_images})\n")
            f.write(f"Average Confidence: {avg_confidence:.4f}\n\n")
            
            # Write per-class statistics
            f.write("Per-Class Statistics:\n")
            for cls_name, stats in class_stats.items():
                if stats['total'] > 0:
                    cls_accuracy = stats['correct'] / stats['total'] * 100
                    cls_avg_conf = stats['confidence_sum'] / stats['correct'] if stats['correct'] > 0 else 0
                    f.write(f"  {cls_name}: Accuracy {cls_accuracy:.2f}% ({stats['correct']}/{stats['total']}), Avg Confidence: {cls_avg_conf:.4f}\n")
            
        else:
            f.write("No valid detections to calculate statistics.\n")
    
    print(f"\nValidation results saved to {os.path.abspath(results_file)}")
    print(f"Result images saved to {os.path.abspath(output_dir)}")

if __name__ == '__main__':
    main()
